# Project RoboOrchard
#
# Copyright (c) 2024-2025 Horizon Robotics. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.

import asyncio
import json
import os
import sys
import threading
import time
import traceback
from pathlib import Path
from typing import Dict
import subprocess
import cv2
import cv_bridge
import numpy as np
import rclpy
print("python location:", sys.executable)
try:
    import websockets
except ImportError:
    print("websockets not found, installing...")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "websockets"])
    import websockets
from message_filters import ApproximateTimeSynchronizer, Subscriber
from rclpy.node import Node, ParameterDescriptor
from sensor_msgs.msg import CameraInfo, Image, JointState

from robo_orchard_deploy.config import DeployConfig


class DeployNodeWS(Node):
    def __init__(self):
        super().__init__("dual_arm_sync_deploy_node_ws")
        self._initialize()
        self.latest_obs = {}
        self.obs_lock = threading.Lock()
        self.left_arm_publisher = self.create_publisher(JointState, self.config.left_arm_control_topic, 1)
        self.right_arm_publisher = self.create_publisher(JointState, self.config.right_arm_control_topic, 1)
        self.observation_subs = []
        for obs in self.config.observation:
            if obs == "left_arm":
                self.left_arm_state_sub = Subscriber(self, JointState, self.config.left_arm_state_topic)
                self.observation_subs.append(self.left_arm_state_sub)
            elif obs == "right_arm":
                self.right_arm_state_sub = Subscriber(self, JointState, self.config.right_arm_state_topic)
                self.observation_subs.append(self.right_arm_state_sub)
            elif obs == "left_rgb":
                self.left_rgb_sub = Subscriber(self, Image, self.config.left_rgb_topic)
                self.observation_subs.append(self.left_rgb_sub)
            elif obs == "left_depth":
                self.left_depth_sub = Subscriber(self, Image, self.config.left_depth_topic)
                self.observation_subs.append(self.left_depth_sub)
            elif obs == "middle_rgb":
                self.middle_rgb_sub = Subscriber(self, Image, self.config.middle_rgb_topic)
                self.observation_subs.append(self.middle_rgb_sub)
            elif obs == "middle_depth":
                self.middle_depth_sub = Subscriber(self, Image, self.config.middle_depth_topic)
                self.observation_subs.append(self.middle_depth_sub)
            elif obs == "right_rgb":
                self.right_rgb_sub = Subscriber(self, Image, self.config.right_rgb_topic)
                self.observation_subs.append(self.right_rgb_sub)
            elif obs == "right_depth":
                self.right_depth_sub = Subscriber(self, Image, self.config.right_depth_topic)
                self.observation_subs.append(self.right_depth_sub)
            elif obs == "left_camera_intrinsic":
                self.left_camera_intrinsic_sub = Subscriber(self, CameraInfo, self.config.left_camera_intrinsic_topic)
                self.observation_subs.append(self.left_camera_intrinsic_sub)
            elif obs == "middle_camera_intrinsic":
                self.middle_camera_intrinsic_sub = Subscriber(self, CameraInfo, self.config.middle_camera_intrinsic_topic)
                self.observation_subs.append(self.middle_camera_intrinsic_sub)
            elif obs == "right_camera_intrinsic":
                self.right_camera_intrinsic_sub = Subscriber(self, CameraInfo, self.config.right_camera_intrinsic_topic)
                self.observation_subs.append(self.right_camera_intrinsic_sub)
            else:
                raise ValueError(f"Unknown observation type: {obs}")
        self.sync = ApproximateTimeSynchronizer(self.observation_subs, queue_size=1, slop=0.05)
        self.sync.registerCallback(self._observe_callback)
        self.cv_bridge = cv_bridge.CvBridge()
        self.declare_parameter("websocket_host", "127.0.0.1")
        self.declare_parameter("websocket_port", 8765)
        self.websocket_host = self.get_parameter("websocket_host").get_parameter_value().string_value
        self.websocket_port = self.get_parameter("websocket_port").get_parameter_value().integer_value
        self.websocket_url = f"ws://{self.websocket_host}:{self.websocket_port}"
        self.ws_thread = threading.Thread(target=self._run_websocket, daemon=True)
        self.ws_thread.start()
        self.get_logger().info(f"Initialized DeployNodeWS with websocket server {self.websocket_url}")

    def _initialize(self):
        self.declare_parameter("config_file", "", descriptor=ParameterDescriptor(description="Config path"))
        config_file: str = self.get_parameter("config_file").get_parameter_value().string_value
        if not os.path.exists(config_file):
            raise FileNotFoundError(f"config file {config_file} does not exists!")
        with open(config_file, "r") as f:
            self.config: DeployConfig = DeployConfig.model_validate_json(f.read())

    def _get_data_keys_info(self, data_dict: dict) -> str:
        if not data_dict:
            return "[]"
        keys_info = {}
        for key, value in data_dict.items():
            if key == "obs" and isinstance(value, dict):
                keys_info[key] = f"dict_keys={list(value.keys())}"
            elif isinstance(value, dict):
                keys_info[key] = f"dict_keys={list(value.keys())}"
            elif isinstance(value, (list, tuple)):
                keys_info[key] = f"list_len={len(value)}"
            else:
                keys_info[key] = type(value).__name__
        return str(keys_info)

    def _run_websocket(self):
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            loop.run_until_complete(self._websocket_client_handler())
        except Exception as e:
            self.get_logger().error(f"WebSocket error: {e}")

    async def _websocket_client_handler(self):
        while True:
            try:
                self.get_logger().info(f"Connecting to websocket server: {self.websocket_url}")
                async with websockets.connect(self.websocket_url, max_size=None) as websocket:
                    self.get_logger().info(f"Connected to websocket server: {self.websocket_url}")
                    async for message in websocket:
                        try:
                            packet_dict = json.loads(message)
                            packet_type = packet_dict.get("type", "")
                            self.get_logger().info(f"Received packet type: {packet_type}")
                            self.get_logger().info(f"Calling _handle_packet for type: {packet_type}")
                            response = await self._handle_packet(packet_type, packet_dict.get("data", {}))
                            data_info = self._get_data_keys_info(response.get('data', {})) if response else "{}"
                            self.get_logger().info(f"_handle_packet returned: type={response.get('type') if response else None}, data_info={data_info}")
                            if response:
                                response_str = json.dumps(response, default=str)
                                data_info = self._get_data_keys_info(response.get('data', {}))
                                self.get_logger().info(f"准备发送响应: type={response.get('type', 'UNKNOWN')}, data_info={data_info}")
                                try:
                                    self.get_logger().info(f"准备发送响应到WebSocket (连接对象: {id(websocket)})")
                                    await websocket.send(response_str)
                                    self.get_logger().info(f"已发送响应: type={response.get('type', 'UNKNOWN')}, 消息长度={len(response_str)}, 连接对象: {id(websocket)}")
                                except Exception as send_error:
                                    self.get_logger().error(f"发送响应时发生错误: {send_error}")
                                    self.get_logger().error(traceback.format_exc())
                                    raise
                            else:
                                self.get_logger().warning(f"No response for packet type: {packet_type}, response is None or empty")
                        except json.JSONDecodeError as e:
                            self.get_logger().error(f"JSON decode error: {e}")
                            error_response = {"type": "ERROR_RESPONSE", "data": {"success": False, "error": f"JSON decode error: {e}"}}
                            await websocket.send(json.dumps(error_response, default=str))
                        except Exception as e:
                            self.get_logger().error(f"Error handling packet: {e}")
                            self.get_logger().error(traceback.format_exc())
                            error_response = {"type": "ERROR_RESPONSE", "data": {"success": False, "error": str(e)}}
                            await websocket.send(json.dumps(error_response, default=str))
            except websockets.exceptions.ConnectionClosed:
                self.get_logger().info("WebSocket connection closed, reconnecting...")
            except ConnectionRefusedError:
                self.get_logger().warning("WebSocket connection refused, retrying in 5 seconds...")
                await asyncio.sleep(5)
            except Exception as e:
                self.get_logger().error(f"WebSocket error: {e}, retrying in 5 seconds...")
                await asyncio.sleep(5)

    async def _handle_packet(self, packet_type: str, data: dict):
        if packet_type == "GET_OBS":
            self.get_logger().info("Received GET_OBS request")
            try:
                obs_dict = self._get_obs_callback()
                serialized = self._serialize_observation(obs_dict)
                self.get_logger().info("Successfully got observation")
                return {"type": "GET_OBS_RESPONSE", "data": {"obs": serialized, "success": True, "error": None}}
            except Exception as e:
                self.get_logger().error(f"Error getting observation: {e}")
                self.get_logger().error(traceback.format_exc())
                return {"type": "GET_OBS_RESPONSE", "data": {"obs": None, "success": False, "error": str(e)}}
        elif packet_type == "TAKE_ACTION":
            self.get_logger().info("Received TAKE_ACTION request")
            try:
                action_data = data.get("action", [])
                self.get_logger().info(f"Action data length: {len(action_data)}")
                action_array = np.array(action_data, dtype=np.float32)
                self.get_logger().info(f"Action array shape: {action_array.shape}")
                result = self._take_action_callback(action_array)
                self.get_logger().info(f"Successfully executed action, result={result}")
                response = {"type": "TAKE_ACTION_RESPONSE", "data": {"success": result.get("success", True), "error": result.get("error", None)}}
                data_info = self._get_data_keys_info(response.get('data', {}))
                self.get_logger().info(f"_handle_packet returning response: type={response.get('type')}, data_info={data_info}")
                return response
            except Exception as e:
                self.get_logger().error(f"Error taking action: {e}")
                self.get_logger().error(traceback.format_exc())
                error_response = {"type": "TAKE_ACTION_RESPONSE", "data": {"success": False, "error": str(e)}}
                data_info = self._get_data_keys_info(error_response.get('data', {}))
                self.get_logger().info(f"_handle_packet returning error response: type={error_response.get('type')}, data_info={data_info}")
                return error_response
        else:
            self.get_logger().warning(f"Unknown packet type: {packet_type}")
            return {"type": "ERROR_RESPONSE", "data": {"success": False, "error": f"Unknown packet type: {packet_type}"}}

    def _serialize_observation(self, obs_dict: Dict) -> Dict:
        serialized = {}
        for key, value in obs_dict.items():
            if isinstance(value, np.ndarray):
                serialized[key] = value.tolist()
            elif isinstance(value, dict):
                serialized[key] = self._serialize_observation(value)
            elif isinstance(value, (list, tuple)):
                serialized[key] = [self._serialize_observation(item) if isinstance(item, dict) else (item.tolist() if isinstance(item, np.ndarray) else item) for item in value]
            else:
                serialized[key] = value
        return serialized

    def _observe_callback(self, *args):
        obs_dict = {}
        left_joint_state = None
        right_joint_state = None
        for obs, msg in zip(self.config.observation, args, strict=False):
            if obs in ["left_rgb", "middle_rgb", "right_rgb"]:
                bgr_img = self.cv_bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
                rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
                obs_dict[obs] = rgb_img
            elif obs in ["left_depth", "middle_depth", "right_depth"]:
                depth_img = self.cv_bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
                obs_dict[obs] = depth_img.astype(np.float32) / 1000.0
            elif obs in ["left_camera_intrinsic", "middle_camera_intrinsic", "right_camera_intrinsic"]:
                intrinsic_matrix = np.array(msg.p).reshape(3, 4)[:, :3]
                obs_dict[obs] = intrinsic_matrix
            elif obs == "left_arm":
                left_joint_state = np.array(msg.position, dtype=np.float32)
            elif obs == "right_arm":
                right_joint_state = np.array(msg.position, dtype=np.float32)
        if left_joint_state is not None and right_joint_state is not None:
            joint_state = np.concatenate([left_joint_state, right_joint_state], axis=0)
            obs_dict["joint_state"] = joint_state
        with self.obs_lock:
            self.latest_obs = obs_dict

    def _get_obs_callback(self) -> Dict:
        with self.obs_lock:
            obs_dict = self.latest_obs.copy()
        result = {}
        middle_rgb = obs_dict.get("middle_rgb", None)
        if middle_rgb is not None:
            if middle_rgb.shape[:2] == (480, 640):
                middle_rgb = middle_rgb[::2, ::2]
            result["image"] = middle_rgb
        else:
            result["image"] = np.zeros((240, 320, 3), dtype=np.uint8)
        middle_depth = obs_dict.get("middle_depth", None)
        if middle_depth is not None:
            if middle_depth.shape[:2] == (480, 640):
                middle_depth = middle_depth[::2, ::2]
            result["depth"] = middle_depth.astype(np.float32)
        else:
            result["depth"] = np.zeros((240, 320), dtype=np.float32)
        point_cloud = obs_dict.get("middle_depth", None)
        if point_cloud is not None:
            result["point_cloud"] = np.zeros((1000, 3), dtype=np.float32)
        else:
            result["point_cloud"] = np.zeros((1000, 3), dtype=np.float32)
        joint_state = obs_dict.get("joint_state", None)
        if joint_state is not None:
            if joint_state.shape[0] != 14:
                self.get_logger().warning(f"Joint state shape mismatch: expected 14, got {joint_state.shape[0]}")
                if joint_state.shape[0] < 14:
                    joint_state = np.pad(joint_state, (0, 14 - joint_state.shape[0]), mode='constant')
                else:
                    joint_state = joint_state[:14]
            result["joint_state"] = joint_state.astype(np.float32)
        else:
            result["joint_state"] = np.zeros(14, dtype=np.float32)
        result["ee_state"] = np.zeros(16, dtype=np.float32)
        return result

    def _take_action_callback(self, action: np.ndarray) -> Dict:
        try:
            self.get_logger().info(f"_take_action_callback called with action shape: {action.shape}")
            assert action.shape == (14,), f"Expected action shape (14,), got {action.shape}"
            left_arm_target = action[0:7].tolist()
            right_arm_target = action[7:14].tolist()
            self.get_logger().info(f"Preparing to send action: left={left_arm_target}, right={right_arm_target}")
            self._send_action(left_arm_target, right_arm_target)
            self.get_logger().info(f"Executed action: left={left_arm_target}, right={right_arm_target}")
            result = {"success": True, "error": None}
            self.get_logger().info(f"_take_action_callback returning: {result}")
            return result
        except Exception as e:
            self.get_logger().error(f"Error executing action: {e}")
            self.get_logger().error(traceback.format_exc())
            error_result = {"success": False, "error": str(e)}
            self.get_logger().info(f"_take_action_callback returning error: {error_result}")
            return error_result

    def _send_left_arm_action(self, joint_position):
        goal_msg = JointState()
        goal_msg.name = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"]
        goal_msg.position = joint_position
        goal_msg.velocity = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 80.0]
        goal_msg.effort = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5]
        self.left_arm_publisher.publish(goal_msg)

    def _send_right_arm_action(self, joint_position):
        goal_msg = JointState()
        goal_msg.name = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"]
        goal_msg.position = joint_position
        goal_msg.velocity = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 80.0]
        goal_msg.effort = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5]
        self.right_arm_publisher.publish(goal_msg)

    def _send_action(self, left_arm_target, right_arm_target):
        self._send_left_arm_action(left_arm_target)
        self._send_right_arm_action(right_arm_target)


def main(args=None):
    rclpy.init(args=args)
    executor = rclpy.executors.MultiThreadedExecutor()
    node = DeployNodeWS()
    executor.add_node(node)
    try:
        executor.spin()
    except KeyboardInterrupt:
        pass
    finally:
        executor.shutdown()
        rclpy.shutdown()

